# *_*coding:utf-8 *_*
# @Author : Reggie
# @Time : 2022/2/17 15:26
import json
import logging
import os
import shlex
import subprocess
import sys

import netifaces
import paramiko

ENVIRON = {
    "OFMT": "JSON",
    "LANG": "en_US.UTF-8;"
}
LOG_LEVEL = "INFO"

env_debug = True if os.environ.get("DEBUG") else False
if env_debug:
    LOG_LEVEL = "DEBUG"

nolog = True if os.environ.get("nolog") else False
if not nolog:
    root_logger = logging.getLogger()
    stream_handle = logging.StreamHandler()
    formatter = logging.Formatter("%(asctime)s -- %(levelname)s --"
                                  " %(message)s (%(funcName)s in %(filename)s %(lineno)d)")
    stream_handle.setFormatter(formatter)
    root_logger.addHandler(stream_handle)
    root_logger.setLevel(LOG_LEVEL)


def run_command(cmd,
                shell=False,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                stdin=subprocess.PIPE,
                input=None,
                cmd_env=None,
                raise_error=False,
                json_out=False,
                encode="utf-8"):
    out, err, rc = {}, "", 1

    if shell:
        if isinstance(cmd, (list, tuple)):
            cmd = " ".join([str(i) for i in cmd])
    else:
        if isinstance(cmd, str):
            cmd = shlex.split(cmd)
        elif isinstance(cmd, (list, tuple)):
            cmd = [str(i) for i in cmd]
    cmd_str = " ".join(cmd) if isinstance(cmd, (list, tuple)) else cmd

    env = dict(os.environ)
    if cmd_env:
        for k in cmd_env.keys():
            env[k] = cmd_env[k]

    try:
        logging.debug(f"Running command: {cmd_str}")
        p = subprocess.Popen(args=cmd, shell=shell, stdout=stdout, stderr=stderr, stdin=stdin, env=env)
        out, err = p.communicate(input=input)
        rc = p.returncode
    except Exception as e:
        err = f"Run Command Error: {e}"
        if raise_error:
            raise Exception(e)
        return out, err, rc

    if encode:
        out = out.decode(encode)
        err = err.decode(encode)

    if json_out:
        try:
            out = json.loads(out)
        except Exception as e:
            err = f"JSON Loads Error: {e}"
            if raise_error:
                raise Exception(e)
            return out, err, rc

    return out, err, rc


class SSHConnection(object):
    def __init__(self, host, port=22, username='root', pwd='user@dev', use_callback=False):
        self.host = host
        self.port = port
        self.username = username
        self.pwd = pwd
        self.use_callback = use_callback
        self.__transport = None
        self._ssh = None
        self._sftp = None

    def __enter__(self):
        if self.__transport is not None:
            raise RuntimeError("Already connected")
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def connect(self):
        if self.__transport is not None and self.__transport.active:
            return
        transport = paramiko.Transport((self.host, self.port))
        transport.connect(username=self.username, password=self.pwd)
        self.__transport = transport

    def close(self):
        self.__transport.close()
        self.__transport = None

    def translate_byte(self, B):
        B = float(B)
        KB = float(1024)
        MB = float(KB ** 2)
        GB = float(MB ** 2)
        TB = float(GB ** 2)
        if B < KB:
            return '{} {}'.format(B, 'bytes' if B > 1 else 'byte')
        elif KB < B < MB:
            return '{:.2f} KB'.format(B / KB)
        elif MB < B < GB:
            return '{:.2f} MB'.format(B / MB)
        elif GB < B < TB:
            return '{:.2f} GB'.format(B / GB)
        else:
            return '{:.2f} TB'.format(B / TB)

    def call_back(self, transferred, toBeTransferred, suffix=''):
        bar_len = 100
        filled_len = int(round(bar_len * transferred / float(toBeTransferred)))
        percents = round(100.0 * transferred / float(toBeTransferred), 1)
        bar = '\033[32;1m%s\033[0m' % '=' * filled_len + '-' * (bar_len - filled_len)
        sys.stdout.write('[%s] %s%s %s\r\n' % (bar, '\033[32;1m%s\033[0m' % percents, '%', suffix))
        sys.stdout.flush()

    def sftp_down_file(self, server_path, local_path):
        if self._sftp:
            sftp = self._sftp
        else:
            sftp = paramiko.SFTPClient.from_transport(self.__transport)
            self._sftp = sftp
        result = sftp.get(server_path, local_path, callback=self.call_back if self.use_callback else None)
        return result

    def sftp_upload_file(self, local_path, server_path):
        if self._sftp:
            sftp = self._sftp
        else:
            sftp = paramiko.SFTPClient.from_transport(self.__transport)
            print("init")
            self._sftp = sftp
        result = sftp.put(local_path, server_path, callback=self.call_back if self.use_callback else None)
        return result

    def cmd(self,
            command,
            encoding="utf-8",
            bufsize=-1,
            timeout=None,
            get_pty=False,
            environment=None,
            ):
        if self._ssh:
            ssh = self._ssh
        else:
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh._transport = self.__transport
            self._ssh = ssh
        # 执行命令
        stdin, stdout, stderr = ssh.exec_command(
            command,
            bufsize=bufsize,
            timeout=timeout,
            get_pty=get_pty,
            environment=environment,
        )
        # 获取命令结果
        stdout = stdout.read().decode(encoding)
        stderr = stderr.read().decode(encoding)
        return stdout, stdin, stderr


def remove_vip(vip) -> (bool, str):
    ipv4_mask2prefix = lambda mask: sum(bin(int(i, 10)).count('1') for i in mask.split('.'))
    netmask = None
    for netface in netifaces.interfaces():
        addrs = netifaces.ifaddresses(netface)
        if netifaces.AF_INET not in addrs:
            continue
        for ip in addrs[netifaces.AF_INET]:
            addr = ip.get("addr")
            if vip == addr:
                netmask = ipv4_mask2prefix(ip["netmask"])
                break
        if netmask:
            break
    else:
        return False, "VIP not found"
    cmdline = f"/usr/sbin/ip add del {vip}/{netmask} dev {netface}"
    out, err, rc = run_command(
        cmdline,
        json_out=False,
    )
    if rc != 0:
        logging.error(f"remove vip failed: out:{out} err:{err} rc:{rc}")
        return False, "remove vip failed"
    return True, "success"


def reset_node(node_ip):
    logging.info(f"reset node: {node_ip}")
    with SSHConnection(node_ip) as client:
        # stdout, stdin, stderr = client.cmd("""systemctl status uxsagent""")
        # if stderr:
        #     logging.error(f"stop all uus service error: {stderr}")
        #     return
        # logging.debug(f"domain force remove: stop all uus service success")
        # stdout, stdin, stderr = client.cmd("""rm -fr  /etc/uraid/*""")
        # if stderr:
        #     logging.error(f"remove /etc/uraid error: {stderr}")
        #     return
        # logging.debug(f"domain force remove: remove /etc/uraid success")
        # stdout, stdin, stderr = client.cmd("""bash /usr/lib/uraid/scripts/clean-up.sh""")
        # if stderr:
        #     logging.error(f"run clean-up.sh error: {stderr}")
        #     return
        # logging.debug(f"domain force remove: run clean-up.sh success")
        stdout, stdin, stderr = client.cmd(
            "ETCDCTL_API=3 etcdctl --endpoints=http://127.0.0.1:22379 get --prefix /uus/domains/domain_info")
        if stderr:
            logging.error(f"get domain_info error: {stderr}")
            return
        domains = stdout.split("\n")
        domains = zip([domain for i, domain in enumerate(domains) if i % 2 == 0],
            [domain for i, domain in enumerate(domains) if i % 2 == 1])
        for etcd_key, etcd_value in domains:
            domain = json.loads(etcd_value)
            if domain.get("node_ip") == node_ip:
                logging.debug(f"find node_vip")
                break
        else:
            logging.error(f"can't find node ip")
            return
        domain_vip = domain.get("domain_vip")
        print(domain_vip)
        # vip = ""
        # remove_vip(vip)


def main(ip=None):
    if ip:
        with SSHConnection(ip) as connect:
            stdout, stdin, stderr = connect.cmd("OFMT=JSON ucli node list")
            if stderr:
                logging.error(f"node list failed: {stderr}")
                return
            out = json.loads(stdout)
    else:
        out, err, rc = run_command("ucli node list", cmd_env={"OFMT": "JSON"})
        if rc != 0:
            logging.error(f"node list failed: {err}")
            return
    node_infos = out.get("data", [])
    for node in node_infos:
        node_ip = node.get("ipv4")
        reset_node(node_ip)


if __name__ == '__main__':
    main("10.200.104.91")
